import numpy as np
import mxnet as mx
from multiprocessing.dummy import Pool
from tqdm import tqdm

path = '/home/wuyuxiang/faces_ms1m_112x112/train.idx'

imgrec = mx.recordio.MXIndexedRecordIO(path[0:-3] + 'idx', path[0:-3] + 'rec', 'r')
header, _ = mx.recordio.unpack(imgrec.read_idx(0))
num_imgs = int(header.label[0]) - 1
num_people = int(header.label[1] - header.label[0])

all_label = num_imgs


def f(data):
    i, imgdata = data
    header, img = mx.recordio.unpack_img(imgdata)
    assert img.shape == (112, 112, 3)
    label = header.label
    img = np.transpose(img[:, :, ::-1], axes=(2, 0, 1))
    np.save(f'/home/wuyuxiang/ms1m/{str(i)}.npy', img)
    return i, label


with Pool(40) as p:
    imgdata = map(lambda x: (x, imgrec.read_idx(x + 1)), range(num_imgs))
    with tqdm(p.imap_unordered(f, imgdata), total=num_imgs) as pbar:
        for i, label in pbar:
            all_label[i] = int(label)

np.save(f'/home/wuyuxiang/ms1m/all_label.npy', all_label)
